import numpy as np  # 导入 NumPy 库，用于数值计算
import plotly.graph_objects as go  # 导入 Plotly 的图形对象库

from util.TDataGenerator import DataGenerator


def load_data(filename):
    data = np.loadtxt(filename, delimiter=',')
    X = data[:, :-1]  # 特征
    y = data[:, -1]  # 目标值
    return X, y


def linear_regression(X, y):
    # 添加偏置项
    X_b = np.c_[np.ones((X.shape[0], 1)), X]  # 在特征矩阵中添加一列1
    # 计算线性方程的参数
    theta_best = np.linalg.inv(X_b.T.dot(X_b)).dot(X_b.T).dot(y)
    return theta_best


def compute_loss(X, y, theta):
    X_b = np.c_[np.ones((X.shape[0], 1)), X]  # 添加偏置项
    predictions = X_b.dot(theta)
    loss = np.mean((predictions - y) ** 2)  # 均方误差
    return loss


def plot_scatter_and_plane(theta):
    # 生成 x1 和 x2 的数据点
    x1 = np.linspace(-10, 10, 100)  # 生成从 -10 到 10 的 100 个等间距点，作为 x1 的数据
    x2 = np.linspace(-10, 10, 100)  # 生成从 -10 到 10 的 100 个等间距点，作为 x2 的数据

    # 创建 x1 和 x2 的二维网格
    x1, x2 = np.meshgrid(x1, x2)  # 使用 meshgrid 函数将 x1 和 x2 转换为二维网格

    # 计算 y 值
    y = theta[0] * x1 + theta[1] * x2 + theta[2]  # 根据方程 y = 0.1 * x1 + 0.3 * x2 计算 y 值

    # 创建 Plotly 图形对象
    fig = go.Figure(data=[go.Surface(z=y, x=x1, y=x2)])  # 创建一个三维表面图，z 轴为 y，x 轴为 x1，y 轴为 x2

    # 生成随机散点数据
    np.random.seed(42)  # 设置随机种子
    random_x1 = np.random.uniform(-10, 10, 50)  # 生成 50 个随机 x1 值
    random_x2 = np.random.uniform(-10, 10, 50)  # 生成 50 个随机 x2 值
    random_y = 0.1 * random_x1 + 0.3 * random_x2  # 根据方程计算对应的 y 值

    # 将随机散点添加到图形中
    fig.add_trace(go.Scatter3d(x=random_x1, y=random_x2, z=random_y, mode='markers', name='随机散点'))


    # 更新图形布局
    fig.update_layout(
        title='3D Surface',  # 设置图形标题
        autosize=False,  # 关闭自动调整大小
        width=700,  # 设置图形宽度为 700 像素
        height=700,  # 设置图形高度为 700 像素
        margin=dict(l=65, r=50, b=65, t=90)  # 设置图形的边距
    )

    # 显示图形
    fig.show()  # 显示图形


if __name__ == "__main__":
    num_rows = 20  # 指定行数
    num_cols = 2  # 指定列数

    generator = DataGenerator(num_rows, num_cols)
    filename = 'train_output/trainNum01.txt'  # 文件名

    # 写入数据到文件
    generator.write_data_to_file(filename)

    # 从文件读取数据
    X, y = load_data(filename)

    # 拟合线性方程
    theta = linear_regression(X, y)

    # 计算损失函数
    loss = compute_loss(X, y, theta)

    # 打印线性方程的参数和损失函数
    print("线性方程的参数:")
    print(theta)
    print(theta[0])
    print(theta[1])
    print(theta[2])
    print("损失函数 (均方误差):")
    print(loss)

    plot_scatter_and_plane(theta)
